package org.chen.config;

import com.alibaba.druid.pool.DruidDataSource;
import com.baomidou.mybatisplus.extension.spring.MybatisSqlSessionFactoryBean;
import org.apache.ibatis.session.SqlSessionFactory;
import org.apache.shardingsphere.driver.api.ShardingSphereDataSourceFactory;
import org.apache.shardingsphere.infra.config.algorithm.AlgorithmConfiguration;
import org.apache.shardingsphere.sharding.api.config.ShardingRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.rule.ShardingTableRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.strategy.sharding.StandardShardingStrategyConfiguration;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.jdbc.datasource.DataSourceTransactionManager;

import javax.sql.DataSource;
import java.sql.SQLException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Properties;

@Configuration
public class ShardingConfig {

    @Value("${spring.shardingsphere.datasource.ds.url}")
    private String url;

    @Value("${spring.shardingsphere.datasource.ds.username}")
    private String username;

    @Value("${spring.shardingsphere.datasource.ds.password}")
    private String password;

    @Primary
    @Bean("shardingDataSource")
    public DataSource shardingDataSource() throws SQLException {
        // 配置数据源
        Map<String, DataSource> dataSourceMap = new HashMap<>();
        DruidDataSource dataSource = new DruidDataSource();
        dataSource.setUrl(url);
        dataSource.setUsername(username);
        dataSource.setPassword(password);
        dataSource.setDriverClassName("com.mysql.cj.jdbc.Driver");
        dataSourceMap.put("ds", dataSource);

        // 配置分片规则
        ShardingRuleConfiguration shardingRuleConfig = new ShardingRuleConfiguration();

        // 配置 comment_task 表规则
        ShardingTableRuleConfiguration tableRuleConfig =
                new ShardingTableRuleConfiguration("comment_task", "ds.comment_task_${0..15}");

        ShardingTableRuleConfiguration tableRuleConfig2 =
                new ShardingTableRuleConfiguration("comment_rule", "ds.comment_rule_${0..10}");

        // 配置分片策略
        tableRuleConfig.setTableShardingStrategy(
                new StandardShardingStrategyConfiguration("id", "comment_task_inline")
        );

        tableRuleConfig2.setTableShardingStrategy(
                new StandardShardingStrategyConfiguration("id", "comment_rule_inline")
        );

        shardingRuleConfig.getTables().add(tableRuleConfig);
        shardingRuleConfig.getTables().add(tableRuleConfig2);

        // 配置分片算法 - 为每个表使用独立的Properties对象
        Properties propsTask = new Properties();
        propsTask.setProperty("algorithm-expression", "comment_task_${id % 16}");
        shardingRuleConfig.getShardingAlgorithms().put(
                "comment_task_inline",
                new AlgorithmConfiguration("INLINE", propsTask)
        );

        Properties propsRule = new Properties();
        propsRule.setProperty("algorithm-expression", "comment_rule_${id % 11}");
        shardingRuleConfig.getShardingAlgorithms().put(
                "comment_rule_inline",
                new AlgorithmConfiguration("INLINE", propsRule)
        );

        // 创建 ShardingSphere 数据源
        Properties properties = new Properties();
        properties.setProperty("sql-show", "true");

        return ShardingSphereDataSourceFactory.createDataSource(
                dataSourceMap,
                Collections.singleton(shardingRuleConfig),
                properties
        );
    }

    @Primary
    @Bean("shardingSqlSessionFactory")
    public SqlSessionFactory shardingSqlSessionFactory(@Qualifier("shardingDataSource") DataSource dataSource) throws Exception {
        MybatisSqlSessionFactoryBean factoryBean = new MybatisSqlSessionFactoryBean();
        factoryBean.setDataSource(dataSource);
        factoryBean.setTypeAliasesPackage("org.chen.model");
        return factoryBean.getObject();
    }

    @Primary
    @Bean("shardingTransactionManager")
    public DataSourceTransactionManager shardingTransactionManager(@Qualifier("shardingDataSource") DataSource dataSource) {
        return new DataSourceTransactionManager(dataSource);
    }
}